from typing import Tuple, List

import cv2
from ril.util import extract_min_rect, show_obj
from ril_env.envs import SpaceWarEnv
import numpy as np

PLAYER_W = [96, 100]
PLAYER_H = [94, 100]
ENEMY_W = [47, 51]

player_x, player_y = None, None


def space_war_rule(frame, act: int):
    global player_x, player_y
    if frame is None or frame.shape[-2:] != (SpaceWarEnv.SCREEN_WIDTH, SpaceWarEnv.SCREEN_HEIGHT):
        return act, None, True
    objects = extract_min_rect(cv2.transpose(frame))  # 查找物体边缘，得到每个物体的最小矩形

    enemies: List[Tuple[float, float]] = []
    for (index, (x, y, w, h)) in enumerate(objects):
        # if ENEMY_W[0] < w < ENEMY_W[1] or (PLAYER_W[0] < w < PLAYER_W[1] and PLAYER_H[0] < h < PLAYER_H[1]):
        #     show_obj(cv2.transpose(frame), (x + w/2, y + h/2), w, h, str(index))
        if PLAYER_W[0] < w < PLAYER_W[1] and PLAYER_H[0] < h < PLAYER_H[1]:
            player_x, player_y = x + w / 2, y + h / 2
        elif ENEMY_W[0] < w < ENEMY_W[1]:
            enemy_x, enemy_y = x + w / 2, y + h / 2
            enemies.append((enemy_x, enemy_y))
    if player_x is None and player_y is None:
        return SpaceWarEnv.SHOOT, None, False

    closest_enemy_x, closest_enemy_y = None, 0
    cant_left, cant_right = False, False
    for (enemy_x, enemy_y) in enemies:
        # print(enemy_x, enemy_y)
        if player_x - (PLAYER_W[1] + ENEMY_W[1]) / 2 < enemy_x < player_x + (PLAYER_W[1] + ENEMY_W[1]) / 2 and \
                closest_enemy_y < enemy_y < player_y:
            closest_enemy_y, closest_enemy_x = enemy_y, enemy_x
        if player_y - PLAYER_H[1] / 2 < enemy_y < player_y + PLAYER_H[1] / 2:
            cant_left = (enemy_x + ENEMY_W[1] / 2 + 5) >= player_x - PLAYER_W[1] / 2
            cant_right = (enemy_x - ENEMY_W[1] / 2 - 5) <= player_x + PLAYER_W[1] / 2

    # print((player_x, player_y), (closest_enemy_x, closest_enemy_y))
    if closest_enemy_x is None or \
            player_x - ENEMY_W[0] / 2 < closest_enemy_x < player_x + ENEMY_W[0] / 2:
        # 在攻击范围内，射它
        return SpaceWarEnv.SHOOT, None, False
    else:
        # 不在攻击范围内，移动到能射它的地方
        if closest_enemy_x < player_x and not cant_left:
            return SpaceWarEnv.MOVE_LEFT, None, False
        elif closest_enemy_x > player_x and not cant_right:
            return SpaceWarEnv.MOVE_RIGHT, None, False
        else:
            return SpaceWarEnv.SHOOT, None, False
